package com.study.chapter05.operator.transform;

import com.study.entity.Event;
import com.study.chapter05.source.ClickSource;
import org.apache.flink.api.common.functions.Partitioner;
import org.apache.flink.api.java.functions.KeySelector;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.source.RichParallelSourceFunction;

/**
 * @Description:
 * @Author: LiuQun
 * @Date: 2022/7/29 22:15
 */
public class PartitionTest {
    public static void main(String[] args) throws Exception {
        //环境
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        //使用自定义的SourceFunction，注意对于自定义的直接继承SourceFunction的，并行度只能设置为1，否则会抛出异常
        DataStreamSource<Event> stream = env.addSource(new ClickSource());

        //1.shuffle随机分区
        // stream.shuffle().print("shuffle").setParallelism(4);

        //2.rebalance轮询分区
        // stream.rebalance().print("rebalance").setParallelism(4);

        //3.rescale重缩放分区
        env.addSource(new RichParallelSourceFunction<Integer>() {
            @Override
            public void run(SourceContext<Integer> sourceContext) throws Exception {
                for (int i = 1; i <= 8; i++) {
                    // 将奇数发送到索引为1的并行子任务
                    // 将偶数发送到索引为0的并行子任务
                    if ((i + 1) % 2 == getRuntimeContext().getIndexOfThisSubtask()) {
                        sourceContext.collect(i + 1);
                    }
                }
            }

            @Override
            public void cancel() {
            }
        }).setParallelism(2)
                // .rescale()
                // .print("rescale")
                .setParallelism(4);

        //4.broadcast广播
        // stream.broadcast().print("broadcast").setParallelism(4);

        //5.global
        // stream.global().print("global").setParallelism(4);

        //6.用户自定义分区
        env.fromElements(1,2,3,4,5,6,7,8)
                .partitionCustom(new Partitioner<Integer>() {
                    @Override
                    public int partition(Integer key, int numPartitions) {
                        return key % 2;
                    }
                }, new KeySelector<Integer, Integer>() {
                    @Override
                    public Integer getKey(Integer value) throws Exception {
                        return value;
                    }
                })
                .print("custom")
                .setParallelism(2);


        env.execute();


    }
}
